import glob
import os
import hashlib
import re
import chardet
import requests
from langchain_community.chat_models import ChatOpenAI
from langchain_community.document_loaders import (
    PyMuPDFLoader,
    PyPDFLoader,
    TextLoader,
    Docx2txtLoader,
    CSVLoader,
    UnstructuredPowerPointLoader,
    DedocFileLoader
)
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI

import config


async def get_embedding(model: str, text: list[str]|str):
    """
    获取文本的向量表示
    """
    if len(text) == 0:
        return [[]]
    client = OpenAI(base_url=config.model["api_base"], api_key=config.model["api_key"])
    res = client.embeddings.create(input=text, model=model).data
    return [x.embedding for x in res]


def str_to_hash(input_string):
    """
    根据传入字符创生成固定hash，并拼接大写首字母
    """
    # 使用 SHA-256 生成哈希值
    hash_object = hashlib.sha256(input_string.encode('utf-8'))
    hex_dig = hash_object.hexdigest()

    # 提取字母部分
    alpha_only = re.sub(r'[^A-Za-z_]', '', hex_dig)

    # 如果需要固定长度，可以截取前几位
    fixed_length_alpha = alpha_only[:15]

    return f"C{fixed_length_alpha}"


async def doc_loader(file_path: str):
    """
    解析文件
    """
    file_ext = os.path.splitext(file_path)[1].lower()
    # 某些ppt转的pdf文件可能会解析失败，换用其他函数
    if file_ext == ".pdf":
        try:
            return PyPDFLoader(file_path).load()
        except Exception as e:
            print(f"Error loading PDF file: {e}")
            res = []
            tmp = PyMuPDFLoader(file_path).lazy_load()
            for item in tmp:
                res.append(item)
            return res
    if file_ext == ".txt":
        return TextLoader(file_path).load()
    if file_ext == ".doc":
        return DedocFileLoader(file_path).load()
    if file_ext == ".docx":
        return Docx2txtLoader(file_path).load()
    if file_ext == ".csv":
        return CSVLoader(file_path).load()
    if file_ext == ".xlsx" or file_ext == ".xls":
        return DedocFileLoader(file_path).load()
    if file_ext == ".pptx":
        return UnstructuredPowerPointLoader(file_path).load()
    with open(file_path, 'rb') as file:
        raw_data = file.read()
        result = chardet.detect(raw_data)
        encoding = result['encoding']

    with open(file_path, 'r', encoding=encoding, errors='ignore') as f:
        return [Document(page_content=f.read(), metadata={"source": file_path})]


def find_files(directory: str, extension=None):
    """
    遍历指定目录及子目录，返回所有文件名称。
    如果指定了extension，则只返回该类型的文件。
    """
    file_list = []

    # 构建glob pattern，支持匹配所有文件或特定扩展名文件
    pattern = f"{directory}/**/*"
    if extension:
        pattern += f".{extension}"

    for filename in glob.iglob(pattern, recursive=True):
        if os.path.isfile(filename):
            file_list.append(filename)

    return file_list


def merge_doc(documents: list[Document]):
    """
    将多个子文档合并成一个大文档
    """
    if len(documents) > 1:
        text = "\n".join([doc.page_content for doc in documents])
        return Document(page_content=text, metadata=documents[0].metadata)
    else:
        return documents[0]


async def get_keywords(query: str, llm: str) -> str:
    """
    提取传入字符串的关键词，关键词用空格拼接
    """
    message = "从问题中提取关键词，并将返回的关键词用空格分隔，没有提取到关键词就返回空字符串，不要造词，否则会对你惩罚。\n----------\n问题：{query}"
    prompt = ChatPromptTemplate.from_messages([("human", message)])
    client = ChatOpenAI(
        model=llm,
        api_key=config.model["api_key"],
        base_url=config.model["api_base"],
        temperature=0.2,
    )
    chat = {"query": RunnablePassthrough()} | prompt | client | StrOutputParser()
    res = chat.invoke(query)
    return res


async def get_rerank_scores(
        query: str,
        contexts: list,
        reranker: str,
) -> list:
    """
    使用reranker对搜索结果进行重排
    """
    _headers= {
        "Authorization": f"Bearer {config.model['api_key']}"
    }
    url = f"{config.model['api_base']}{config.model[reranker]}"
    response = requests.post(url=url, json={"query": query, "contexts": contexts}, headers=_headers)
    if response.status_code == 200:
        response_json = response.json()
        return response_json["data"]
    else:
        raise Exception(response.text)

